program counterplot
	version 14, missing

// Define the correct syntax and create macros from user input
    syntax varlist(min=3 max=3 numeric) [if] [fweight] [, measure(string) outcome(string)]
    tokenize `varlist'
	args disease exposure strata

// Preliminary checking

* Error if no observations exist
marksample use

quietly count if `exposure' & `use'
	if (r(N)==0) {
	error 2000
	}

quietly count if `disease' & `use'
	if (r(N)==0) {
	error 2000
	}

quietly count if `strata' & `use'
	if (r(N)==0) {
	error 2000
	}

* Error if variables are not dichotomous
tempname N1 N2 N3
quietly tab `exposure' if `use'
scalar `N1' = r(r)

quietly tab `disease' if `use'
scalar `N2' = r(r)

quietly tab `strata' if `use'
scalar `N3' = r(r)

if `N1' + `N2' + `N3' != 6 {
	display as error "Disease, exposure and strata variables must have two categories each"
    exit 498
}

// Get the 2x2 tables numbers

tempname p1_1 p2_1 p1_2 p2_2 s1 s2 z1 z2
tempvar gfplot_dummy_var

quietly egen `gfplot_dummy_var' = group(`strata')

* Outcome Success x Failure

if "`outcome'" == "success" {
	tempvar success
	local title = "Success"
	quietly gen `success' = 1 - `disease'

quietly tab `success' `exposure' if `use' & `gfplot_dummy_var' == 1 [`weight' `exp'], matcell(x)
scalar `p1_1' = (x[2,1] / (x[1,1] + x[2,1]))
scalar `p2_1' = (x[2,2] / (x[1,2] + x[2,2]))

quietly tab `success' `exposure' if `use' & `gfplot_dummy_var' == 2 [`weight' `exp'], matcell(x)
scalar `p1_2' = (x[2,1] / (x[1,1] + x[2,1]))
scalar `p2_2' = (x[2,2] / (x[1,2] + x[2,2]))

quietly tab `exposure' `strata' if `use' [`weight' `exp'], matcell(s)
scalar `s1' = s[1,1] / ((s[1,1] + s[1,2]))
scalar `s2' = s[2,1] / ((s[2,1] + s[2,2]))

quietly tab `success' `exposure' if `use' [`weight' `exp'], matcell(z)
scalar `z1' = z[2,2] / ((z[1,2] + z[2,2]))
scalar `z2' = z[2,1] / ((z[2,1] + z[1,1]))

matrix drop x s z
}

else if "`outcome'" == "failure" | "`outcome'" == "" {

local title = "Failure"

quietly tab `disease' `exposure' if `use' & `gfplot_dummy_var' == 1 [`weight' `exp'], matcell(x)
scalar `p1_1' = (x[2,1] / (x[1,1] + x[2,1]))
scalar `p2_1' = (x[2,2] / (x[1,2] + x[2,2]))

quietly tab `disease' `exposure' if `use' & `gfplot_dummy_var' == 2 [`weight' `exp'], matcell(x)
scalar `p1_2' = (x[2,1] / (x[1,1] + x[2,1]))
scalar `p2_2' = (x[2,2] / (x[1,2] + x[2,2]))

quietly tab `exposure' `strata' if `use' [`weight' `exp'], matcell(s)
scalar `s1' = s[1,1] / ((s[1,1] + s[1,2]))
scalar `s2' = s[2,1] / ((s[2,1] + s[2,2]))

quietly tab `disease' `exposure' if `use' [`weight' `exp'], matcell(z)
scalar `z1' = z[2,2] / ((z[1,2] + z[2,2]))
scalar `z2' = z[2,1] / ((z[2,1] + z[1,1]))

matrix drop x s z
}

else  {
	display as error "Invalid outcome option. Use outcome(failure) for rate of failures; outcome(success) for rate of success"
    exit 498
}

// Select the right graph to display

if "`measure'" == "rd" | "`measure'" == "" {

// Rate differences
twoway function non_exposed = (`p1_1' * x) + (`p1_2' * (1 - x)), ra(0 1) lcolor(red) ///
xtitle(Weight) ytitle(Estimate) yvarlab("not exposed") title("Counterfactual Plot for Rate Differences") subtitle(`title') ///
|| function exposed = (`p2_1' * x) + (`p2_2' * (1 - x)), ra(0 1) lcolor(blue) ///
xline(`=`s1'', lcolor(red)) xline(`=`s2'', lcolor(blue)) yline(`=`z1'', lcolor(blue)) yline(`=`z2'', lcolor(red))
}

else if "`measure'" == "rr" {

// Rate ratios
twoway function non_exposed = (`p1_1' * x) + (`p1_2' * (1 - x)), ra(0 1) lcolor(red) ///
xtitle(Weight) ytitle(Estimate) yvarlab("not exposed") title("Counterfactual Plot for Rate Ratios") subtitle(`title') ///
yscale(log)|| function exposed = (`p2_1' * x) + (`p2_2' * (1 - x)), ra(0 1) lcolor(blue) ///
xline(`=`s1'', lcolor(red)) xline(`=`s2'', lcolor(blue)) yline(`=`z1'', lcolor(blue)) yline(`=`z2'', lcolor(red))
}

else if "`measure'" == "or" {

// LogOdds ratio
twoway function non_exposed = log(((`p1_1' * x) + (`p1_2' * (1 - x))) / (1 - ((`p1_1' * x) + (`p1_2' * (1 - x))))), ra(0 1) lcolor(red) ///
xtitle(Weight) ytitle(Estimate) yvarlab("not exposed") title("Counterfactual Plot for Log Odds Ratios") subtitle(`title') ///
|| function exposed = log(((`p2_1' * x) + (`p2_2' * (1 - x))) / (1 - ((`p2_1' * x) + (`p2_2' * (1 - x))))), ra(0 1) lcolor(blue) ///
xline(`=`s1'', lcolor(red)) xline(`=`s2'', lcolor(blue)) yline(`=log(`z1'/(1-`z1'))', lcolor(blue)) yline(`=log(`z2'/(1-`z2'))', lcolor(red))
}

// Error when no valid measure of association was selected
else  {
	display as error "Invalid measure of association. Use measure(rr) for rate ratios; measure(rd) for rate differences; measure(or) for log odds ratios)"
    exit 498
	}

end